import os
import heapq

from vllm.logger import logger
from vllm.v1.core.sched.request_queue import PriorityRequestQueue


def add_request(self, request) -> None:
    """Add a request to the queue according to priority policy. """
    heapq.heappush(self._heap,
                    (request.num_prompt_tokens, request.arrival_time, request))

batch_by_request_length = os.environ.get("BATCH_BY_REQUEST_LENGTH", "false")
if batch_by_request_length == "true":
    logger.info(f"batch by request length: {batch_by_request_length}")
    PriorityRequestQueue.add_request = add_request
